
import claripy

from ..shellcode import Shellcode


class X86LeakAddress(Shellcode):

    os = ["cgc"]
    arches = ["X86"]
    name = "leakaddress"
    code = [b"\x31\xc0\x40\x40\x31\xdb\x43\xb9", b"\xba", b"\x31\xf6\xcd\x80"]

    def __init__(self, address, length):
        """
        :param addr:    Address to leak.
        :param length:  Amount of bytes to write from the leak.
        """
        self.address = address
        self.length = length

    def raw(self, arch=None):

        address, length = self.address, self.length

        if isinstance(self.address, int):
            address = claripy.BVV(address, 32)
        if isinstance(length, int):
            length = claripy.BVV(length, 32)

        bv_codes = [claripy.BVV(x) for x in self.code]

        return claripy.Concat(bv_codes[0], address.reversed, bv_codes[1], length.reversed, bv_codes[2])
